from abc import ABC, abstractmethod
from typing import Dict, Iterable, Optional, Tuple

from aicenter.models import AIIdentity, AIInstance, AIModel, AIRole
from ._client import AIClient
from ._error import AIIdentityRegistrationError
from ._model import AIModelRequest, AIModelResponse, AITokenUsage, ContentType, MessageRole


class AIChatContent:
    """
    AI 对话内容类
    """

    def __init__(self, _type: ContentType, content: str):
        self._type = _type
        self._content = content

    @property
    def type(self) -> ContentType:
        """AI 对话内容类型"""
        return self._type

    @property
    def content(self) -> str:
        """AI 对话内容"""
        return self._content


class AIChatMessage:
    """
    AI 对话消息
    """

    def __init__(self, role: MessageRole, speaker: str, contents: Iterable[AIChatContent]):
        self._role = role
        self._speaker = speaker
        self._contents = tuple(contents)

    @property
    def role(self) -> MessageRole:
        """消息来源角色"""
        return self._role

    @property
    def speaker(self) -> str:
        """发言者名称"""
        return self._speaker

    @property
    def contents(self) -> Tuple[AIChatContent]:
        """消息内容"""
        return self._contents[:]

    def content(self, index: int) -> AIChatContent:
        """
        获取指定的内容
        :param index: 索引
        :return: 返回对应位置的内容
        """
        return self._contents[index]


class AIChatRequest(AIModelRequest):
    """
    AI 对话请求类
    """

    def __init__(
            self, model: AIModel, role: Optional[AIRole], identity: Optional[AIIdentity],
            history: Optional[Iterable[AIChatMessage]], messages: Iterable[AIChatMessage],
            params: Optional[dict]
    ):
        super().__init__(model=model, params=params)
        self._role: Optional[AIRole] = role
        self._identity: Optional[AIIdentity] = identity
        self._history = tuple(history) if history is not None else tuple()
        self._messages = tuple(messages)

    @property
    def role(self) -> Optional[AIRole]:
        """AI 角色"""
        return self._role

    @property
    def identity(self) -> Optional[AIIdentity]:
        """AI 身份"""
        return self._identity

    @property
    def history(self) -> Tuple[AIChatMessage, ...]:
        """历史对话"""
        return self._history

    @property
    def messages(self) -> Tuple[AIChatMessage, ...]:
        """本次发送的消息"""
        return self._messages


class AIChatResponse(AIModelResponse[AIChatRequest]):
    """
    AI 对话响应类
    """

    def __init__(self, request: AIChatRequest, message: Iterable[AIChatMessage], token_usage: AITokenUsage):
        super().__init__(request=request, token_usage=token_usage)
        self._history = request.history
        self._message = tuple(message)

    @property
    def role(self) -> Optional[AIRole]:
        """AI 角色"""
        return self._request.role

    @property
    def identity(self) -> Optional[AIIdentity]:
        """AI 身份"""
        return self._request.identity

    @property
    def history(self) -> Tuple[AIChatMessage, ...]:
        """历史对话"""
        return self._history

    @property
    def request_message(self) -> Tuple[AIChatMessage, ...]:
        """请求消息"""
        return self._request.messages

    @property
    def response_message(self) -> Tuple[AIChatMessage, ...]:
        """响应消息"""
        return self._message


class AIChatClient(AIClient, ABC):
    """
    对话 AI 客户端
    """

    def __init__(self, instance: AIInstance, *models: AIModel):
        super().__init__(instance, *models)
        self._identities: Dict[int, Dict[int, Tuple[AIRole, AIIdentity]]] = dict()

    def chat(
            self, model_id: int, identity: Optional[Tuple[AIRole, AIIdentity]],
            history: Optional[Iterable[AIChatMessage]], messages: Iterable[AIChatMessage],
            params: Optional[dict] = None
    ) -> AIChatResponse:
        """
        进行一次对话

        :param model_id: AI 模型 ID
        :param identity: AI 角色与身份
        :param history: 历史消息
        :param messages: 消息集
        :param params: 模型参数
        :return: 返回单次对话的结果
        """
        # 获取对应的 AI 模型
        model = self._models[model_id]
        # 声明 AI 角色与 AI 身份
        _role, _identity = None, None
        if identity is not None:
            _role, _identity = identity
            # 检查当前 AI 角色与身份是否已完成注册
            if _identity.id not in self._identities.get(model.id, {}):
                # 注册该 AI 角色与 AI 身份
                self.register_identity(model_id, _role, _identity)
        # 构建 AI Chat 请求对象
        request = AIChatRequest(model, _role, _identity, history, messages, params)
        # 请求 AI 接口，获取响应结果
        messages, token_usages = self._chat(request)
        # 构建响应对象
        return AIChatResponse(request, messages, token_usages)

    def register_identity(self, model_id: int, role: AIRole, identity: AIIdentity) -> None:
        """
        注册 AI 身份

        若注册失败，将抛出 AIIdentityRegistrationError 异常

        :param model_id: AI 模型 ID，必须是已经注册到该 AI 客户端中的模型 ID
        :param role: AI 角色
        :param identity: AI 身份
        """
        # 获取对应的 AI 模型
        model = self._models[model_id]
        # 调用子类的注册函数，若子类注册模型失败，则抛出错误
        if not self._register_identity(model, role, identity):
            raise AIIdentityRegistrationError(model, role, identity)
        # 若注册成功，则绑定该模型
        self._identities[model.id][identity.id] = (role, identity)

    @abstractmethod
    def _register_identity(self, model: AIModel, role: AIRole, identity: AIIdentity) -> bool:
        """
        注册 AI 身份

        :param model: AI 模型
        :param role: AI 角色
        :param identity: AI 身份
        :return: 若返回 False，则表示注册失败
        """
        pass

    @abstractmethod
    def _chat(self, request: AIChatRequest) -> (Tuple[AIChatMessage], AITokenUsage):
        """
        对话实现函数

        :param request: 对话请求参数
        :return: 返回响应消息、令牌使用量
        """
        pass
